import numpy as np
import scipy.signal as signal
import neurokit2 as nk

def check_and_invert_ecg(ecg_signal, sampling_rate=500):
    """
    Checks if the ECG signal is inverted and inverts it if necessary.

    Parameters:
    ecg_signal (numpy.ndarray): ECG signal from a single lead.
    sampling_rate (int): Sampling rate of the ECG signal.

    Returns:
    tuple:
        numpy.ndarray: ECG signal, inverted if necessary.
        bool: True if the signal was inverted, False otherwise.
    """
    ecg_corrected, is_inverted = nk.ecg_invert(ecg_signal, sampling_rate=sampling_rate, show=False)
    return ecg_corrected, is_inverted

def detect_r_peaks_combined(ecg_signal, sampling_rate=500, max_attempts=20):
    """
    Detect R peaks in the ECG signal using scipy's find_peaks,
    adjusting the height threshold dynamically if no peaks are found,
    and enhancing R peaks using the derivative of the signal.

    Parameters:
    ecg_signal (numpy.ndarray): ECG signal from a single lead.
    sampling_rate (int): Sampling rate of the ECG signal.
    max_attempts (int): Maximum number of attempts to detect R peaks with adjusted thresholds.

    Returns:
    numpy.ndarray: Indices of detected R peaks.
    """
    min_distance = int(0.4 * sampling_rate)  # 400 ms to avoid T waves

    # Enhance R peaks by multiplying the signal with its absolute derivative
    diff_signal = np.abs(np.diff(ecg_signal, prepend=ecg_signal[0]))
    enhanced_signal = ecg_signal * diff_signal

    # Normalize the enhanced signal
    enhanced_signal = (enhanced_signal - np.mean(enhanced_signal)) / np.std(enhanced_signal)

    # Initial height threshold using the 99th percentile
    initial_height_threshold = 0.5 * np.percentile(enhanced_signal, 99)

    # Fallback to 10% of max if percentile is too low
    if initial_height_threshold < 0.1 * np.max(enhanced_signal):
        initial_height_threshold = 0.1 * np.max(enhanced_signal)

    height_threshold = initial_height_threshold

    for attempt in range(max_attempts):
        r_peaks_scipy, _ = signal.find_peaks(
            enhanced_signal,
            distance=min_distance,
            height=height_threshold
        )

        if r_peaks_scipy is not None and len(r_peaks_scipy) >= 4:
            # Adjust peaks to local maxima in the original signal
            adjusted_r_peaks = []
            refractory_period = int(0.4 * sampling_rate)  # 400 ms refractory period
            last_peak = -np.inf

            for rp in sorted(r_peaks_scipy):
                if rp - last_peak >= refractory_period:
                    window_size = int(0.1 * sampling_rate)  # 100 ms window
                    window_start = int(max(rp - window_size // 2, 0))
                    window_end = int(min(rp + window_size // 2, len(ecg_signal)))
                    window_signal = ecg_signal[window_start:window_end]

                    # Find local maximum within the window
                    local_max_idx = np.argmax(window_signal)
                    adjusted_peak = window_start + local_max_idx
                    adjusted_r_peaks.append(adjusted_peak)
                    last_peak = adjusted_peak
                else:
                    # Skip peaks within the refractory period
                    continue

            # Remove duplicates and sort
            adjusted_r_peaks = np.unique(adjusted_r_peaks)
            return adjusted_r_peaks
        else:
            # Lower the threshold and retry
            height_threshold *= 0.7  # Reduce threshold for the next attempt
            continue

    # If R peaks are still not detected after all attempts, return an empty array
    return np.array([])

# def detect_r_peaks_combined(ecg_signal, sampling_rate=500):
#     """
#     Detect R peaks in the ECG signal using scipy's find_peaks,
#     adjusting the height threshold to handle noisy signals.

#     Parameters:
#     ecg_signal (numpy.ndarray): ECG signal from a single lead.
#     sampling_rate (int): Sampling rate of the ECG signal.

#     Returns:
#     numpy.ndarray: Indices of detected R peaks.
#     """
#     # Use scipy's find_peaks
#     min_distance = int(0.4 * sampling_rate)  # 400 ms to avoid T waves

#     # Use the 99th percentile to set the height threshold
#     height_threshold = 0.5 * np.percentile(ecg_signal, 99)

#     # If the 99th percentile is too low (e.g., in flat signals), fall back to max value
#     if height_threshold < 0.1 * np.max(ecg_signal):
#         height_threshold = 0.1 * np.max(ecg_signal)

#     r_peaks_scipy, _ = signal.find_peaks(
#         ecg_signal,
#         distance=min_distance,
#         height=height_threshold
#     )

#     # Adjust peaks to local maxima
#     adjusted_r_peaks = []
#     refractory_period = int(0.4 * sampling_rate)  # 400 ms refractory period to skip T waves
#     last_peak = -np.inf

#     for rp in sorted(r_peaks_scipy):
#         if rp - last_peak >= refractory_period:
#             window_size = int(0.4 * sampling_rate)  # 400 ms window
#             window_start = int(max(rp - window_size // 2, 0))
#             window_end = int(min(rp + window_size // 2, len(ecg_signal)))
#             window_signal = ecg_signal[window_start:window_end]

#             local_max_idx = np.argmax(window_signal)
#             adjusted_peak = window_start + local_max_idx
#             adjusted_r_peaks.append(adjusted_peak)
#             last_peak = adjusted_peak
#         else:
#             continue

#     # Remove duplicates and sort
#     adjusted_r_peaks = np.unique(adjusted_r_peaks)

#     return adjusted_r_peaks

def detect_ed_points(ecg_signal, r_peaks_indices, sampling_rate=500):
    """
    Detect ED, ES points (QRS onset) in the ECG signal.

    Parameters:
    ecg_signal (numpy.ndarray): ECG signal from a single lead.
    r_peaks_indices (numpy.ndarray): Indices of R peaks.
    sampling_rate (int): Sampling rate of the ECG signal.

    Returns:
    numpy.ndarray: Indices of detected ED points.
    numpy.ndarray: Indices of detected ES points.
    """
    _, waves = nk.ecg_delineate(ecg_signal, r_peaks_indices, sampling_rate=sampling_rate, method='dwt', show=False)
    if 'ECG_T_Offsets' in waves and 'ECG_P_Offsets' in waves:
        es_points = waves['ECG_T_Offsets']
        ed_points = waves['ECG_P_Offsets']

    ed_points = np.array(ed_points)
    ed_points = ed_points[~np.isnan(ed_points)].astype(int)
    
    es_points = np.array(es_points)
    es_points = es_points[~np.isnan(es_points)].astype(int)
    
    return ed_points, es_points

def detect_ed_points_in_ecg(ecg_data, sampling_rate=500):
    """
    Detect ED points in the ECG data by iterating over leads until R peaks are detected.

    Parameters:
    ecg_data (numpy.ndarray): ECG data array of shape (12, N).
    sampling_rate (int): Sampling rate of the ECG signal.

    Returns:
    dict: Dictionary containing:
        - 'lead_used': Index of the lead used for detection.
        - 'ecg_signal': ECG signal from the selected lead.
        - 'r_peaks_indices': Indices of detected R peaks.
        - 'ed_points': Indices of detected ED points.
    """
    num_leads = ecg_data.shape[0]
    for lead_idx in range(num_leads):
        ecg_signal = ecg_data[lead_idx]

        # Invert the signal if necessary
        ecg_corrected, _ = check_and_invert_ecg(ecg_signal, sampling_rate=sampling_rate)

        # Detect R peaks
        r_peaks_indices = detect_r_peaks_combined(ecg_corrected, sampling_rate=sampling_rate)

        # If R peaks are detected, proceed with this lead
        if len(r_peaks_indices) > 4:  # Arbitrary threshold for sufficient peaks
            # Detect ED points (QRS onset)
            ed_points, es_points = detect_ed_points(ecg_corrected, r_peaks_indices, sampling_rate=sampling_rate)

            return {
                'lead_used': lead_idx,
                'ecg_signal': ecg_signal,  # Original signal (not inverted)
                'r_peaks_indices': r_peaks_indices,
                'ed_points': ed_points,
                'es_points': es_points,
            }
        else:
            continue

    # If R peaks are not detected in any lead, raise an exception or handle accordingly
    raise ValueError("R peaks could not be detected in any lead.")

